import json
import time
from typing import Optional

from .base import BaseNodes, tool_node
from ...mcps import MCPClientManager
from ...models.graph import GraphState, AgentStepLog, GraphRequest
from ...config.agents import AgentActions, SELECTED_MCP_SERVERS
from ...utils.str import unescape_html, print_message
from ...parser import JSONObjOutputParser
from ...prompts.template import get_prompt_template

from langchain_core.runnables import RunnableConfig
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.types import StreamWriter, interrupt
from pydantic import BaseModel
from loguru import logger


class UseMCPToolSchema(BaseModel):
    server_name: str = None
    tool_name: str = None
    arguments: dict = {}


class MCPNodes(BaseNodes):
    def __init__(self, mcp_client_manager: MCPClientManager, **kwargs):
        super().__init__(**kwargs)
        self.mcp_client_manager = mcp_client_manager

    @tool_node(args_schema=UseMCPToolSchema, stream=False, execution_mode="server")
    async def use_mcp_tool(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        调用MCP服务器工具。使用已连接MCP服务器提供的各种功能工具。
        参数:
        - server_name: (必填) MCP服务器名称
        - tool_name: (必填) 工具名称
        - arguments: (必填) 工具参数JSON对象
        用法:
        <use_mcp_tool>
        <server_name>cursor</server_name>
        <tool_name>read_file</tool_name>
        <arguments>{"path": "src/app.js"}</arguments>
        </use_mcp_tool>
        """
        server_name = state.goto.params.get("server_name")
        tool_name = state.goto.params.get("tool_name")
        arguments = state.goto.params.get("arguments", "{}")
        logger.info(
            f"execute use_mcp_tool server_name: {server_name}, tool_name: {tool_name}, arguments: {arguments[:100]}...")
        if not server_name or not tool_name:
            msg = "use_mcp_tool requires server name, tool name"
            logger.error(msg)
            return self.node_failure(msg)

        if self.mcp_client_manager.is_client_type(server_name):
            writer(AgentStepLog.build_tool_action(action=AgentActions.USE_MCP_TOOL.value,
                                                  output={"type": AgentActions.USE_MCP_TOOL.value,
                                                          "params": {
                                                              "server_name": server_name,
                                                              "tool_name": tool_name,
                                                              "arguments": arguments if isinstance(arguments,
                                                                                                   str) else json.dumps(
                                                                  arguments, ensure_ascii=False)
                                                          },
                                                          "kwargs": {}},
                                                  meta={"finish": True}))
            value = interrupt("interrupt")
            return self.node_success(message=value)

        try:
            if isinstance(arguments, str):
                arguments = json.loads(arguments)

            def unescape_arguments(args):
                if isinstance(args, str):
                    return unescape_html(args)
                else:
                    return args

            arguments = unescape_arguments(arguments)

            result = await self.mcp_client_manager.execute_tool(server_name, tool_name, arguments)
            content = result.content[0]
            if result.isError:
                msg = f"""
Error executing use_mcp_tool Result:
<error>
        {content.text}
</error>
"""
                return self.node_failure(msg)
            return self.node_success(f"""
use_mcp_tool Result:
<result>
    {content.text}
</result>
""")
        except Exception as e:
            msg = f"Error while executing use_mcp_tool: {str(e)}"
            return self.node_failure(msg)

    @tool_node()
    async def select_mcp_tools(self, state: GraphState, *, config: Optional[RunnableConfig] = None):
        """智能选择MCP工具节点"""
        question = self._extract_user_question(state.messages)
        if not question:
            return self.node_success()

        request = GraphRequest.from_runnable_config(config)

        try:
            result = await self._select_mcp_tools_by_question(question, state, verbose=request.verbose)

            if not state.ext:
                state.ext = {}
            state.ext[SELECTED_MCP_SERVERS] = result.get("selected_servers", {})
            return self.node_success(ext=state.ext)
        except Exception as e:
            logger.error(f"Error selecting MCP tools: {str(e)}")
            return self.node_success()

    async def _select_mcp_tools_by_question(self, user_question: str, state: GraphState, verbose: bool = False) -> dict:
        default_result = {"need_mcp_tools": False}

        if not self.mcp_client_manager:
            logger.warning("MCP client manager is not available")
            return default_result

        try:
            if not self.mcp_client_manager.is_available():
                await self.mcp_client_manager.initialize()

            mcp_tools_str = await self.mcp_client_manager.get_all_tools_str()
            if not mcp_tools_str:
                return default_result

            system_template = get_prompt_template(
                "tool_mcp_selector",
                mcp_tools=mcp_tools_str,
                user_question=user_question
            )

            messages = [SystemMessage(content=system_template),
                        HumanMessage(content=f"用户当前的问题是：{user_question}")]

            if verbose:
                print_message(messages)

            chain = self.llm_chat.chat_model | JSONObjOutputParser()

            start_time = time.time()
            result, usage = chain.invoke(messages)
            elapsed_time = time.time() - start_time
            logger.info(f"MCP tool selection completed in {elapsed_time:.3f}s, use model: {self.llm_chat.model_name}")

            if usage:
                state.add_usage_meta(self.llm_chat.model_name, usage)

            if not isinstance(result, dict) or not result.get("need_mcp_tools", False):
                return default_result

            available_servers = set(self.mcp_client_manager.connections.keys())
            validated_servers = {
                server: tools for server, tools in result.get("selected_servers", {}).items()
                if server in available_servers
            }

            result["selected_servers"] = validated_servers or {}
            return result
        except Exception as e:
            logger.error(f"Error selecting MCP tools: {str(e)}")
            return default_result
